"""
基于numpy的股票交易环境实现
"""

from __future__ import annotations

import gymnasium as gym
import numpy as np
from numpy import random as rd


class StockTradingEnv(gym.Env):
    """
    股票交易环境类，继承自gym.Env
    该环境模拟股票交易过程，支持多个股票同时交易
    """

    def __init__(
        self,
        config,
        initial_account=1e6,
        gamma=0.99,
        turbulence_thresh=99,
        min_stock_rate=0.1,
        max_stock=1e2,
        initial_capital=1e6,
        buy_cost_pct=1e-3,
        sell_cost_pct=1e-3,
        reward_scaling=2**-11,
        initial_stocks=None,
    ):
        """
        初始化股票交易环境

        参数:
        config: 配置字典，包含价格、技术指标和波动率数据
        initial_account: 初始账户资金 (默认: 1,000,000)
        gamma: 折扣因子，用于计算累积奖励 (默认: 0.99)
        turbulence_thresh: 波动率阈值，超过该值会触发特殊处理 (默认: 99)
        min_stock_rate: 最小股票交易比例 (默认: 0.1)
        max_stock: 最大股票持有数量 (默认: 100)
        initial_capital: 初始资本 (默认: 1,000,000)
        buy_cost_pct: 买入成本比例 (默认: 0.001, 即0.1%)
        sell_cost_pct: 卖出成本比例 (默认: 0.001, 即0.1%)
        reward_scaling: 奖励缩放因子 (默认: 2**-11)
        initial_stocks: 初始股票持仓数组 (默认: None, 表示无持仓)
        """
        # 从配置中提取价格、技术指标和波动率数组
        price_ary = config["price_array"]
        tech_ary = config["tech_array"]
        turbulence_ary = config["turbulence_array"]
        if_train = config["if_train"]
        
        # 将价格和技术指标数组转换为float32类型以提高性能
        self.price_ary = price_ary.astype(np.float32)
        self.tech_ary = tech_ary.astype(np.float32)
        self.turbulence_ary = turbulence_ary

        # 对技术指标进行缩放处理
        self.tech_ary = self.tech_ary * 2**-7
        # 创建波动率布尔数组：超过阈值为1，否则为0
        self.turbulence_bool = (turbulence_ary > turbulence_thresh).astype(np.float32)
        # 对波动率数组应用sigmoid函数并缩放
        self.turbulence_ary = (
            self.sigmoid_sign(turbulence_ary, turbulence_thresh) * 2**-5
        ).astype(np.float32)

        # 获取股票维度（股票数量）
        stock_dim = self.price_ary.shape[1]
        self.gamma = gamma
        self.max_stock = max_stock
        self.min_stock_rate = min_stock_rate
        self.buy_cost_pct = buy_cost_pct
        self.sell_cost_pct = sell_cost_pct
        self.reward_scaling = reward_scaling
        self.initial_capital = initial_capital
        # 初始化股票持仓，如果未提供则初始化为零数组
        self.initial_stocks = (
            np.zeros(stock_dim, dtype=np.float32)
            if initial_stocks is None
            else initial_stocks
        )

        # 环境状态变量，将在reset方法中初始化
        self.day = None  # 当前交易日
        self.amount = None  # 现金金额
        self.stocks = None  # 股票持仓数量
        self.total_asset = None  # 总资产
        self.gamma_reward = None  # 折扣奖励
        self.initial_total_asset = None  # 初始总资产

        # 环境基本信息
        self.env_name = "StockEnv"  # 环境名称
        # 状态维度计算：现金 + (波动率, 波动率布尔值) + (价格, 股票数量, 冷却时间) * 股票数 + 技术指标维度
        self.state_dim = 1 + 2 + 3 * stock_dim + self.tech_ary.shape[1]
        self.stocks_cd = None  # 股票冷却时间
        self.action_dim = stock_dim  # 动作维度（等于股票数量）
        self.max_step = self.price_ary.shape[0] - 1  # 最大步数（交易日数-1）
        self.if_train = if_train  # 是否为训练模式
        self.if_discrete = False  # 是否为离散动作空间
        self.target_return = 10.0  # 目标回报
        self.episode_return = 0.0  # 回合回报

        # 定义观察空间和动作空间
        self.observation_space = gym.spaces.Box(
            low=-3000, high=3000, shape=(self.state_dim,), dtype=np.float32
        )
        self.action_space = gym.spaces.Box(
            low=-1, high=1, shape=(self.action_dim,), dtype=np.float32
        )

    def reset(
        self,
        *,
        seed=None,
        options=None,
    ):
        """
        重置环境到初始状态

        参数:
        seed: 随机种子
        options: 其他选项

        返回:
        state: 初始状态
        info: 额外信息字典
        """
        self.day = 0  # 重置到第一个交易日
        price = self.price_ary[self.day]  # 获取当日股价

        if self.if_train:
            # 训练模式下，为初始股票持仓添加随机扰动
            self.stocks = (
                self.initial_stocks + rd.randint(0, 64, size=self.initial_stocks.shape)
            ).astype(np.float32)
            # 初始化股票冷却时间为0
            self.stocks_cool_down = np.zeros_like(self.stocks)
            # 随机调整初始资金（±5%）
            self.amount = (
                self.initial_capital * rd.uniform(0.95, 1.05)
                - (self.stocks * price).sum()
            )
        else:
            # 测试模式下，使用固定的初始股票持仓
            self.stocks = self.initial_stocks.astype(np.float32)
            self.stocks_cool_down = np.zeros_like(self.stocks)
            self.amount = self.initial_capital  # 使用固定初始资金

        # 计算总资产（现金+股票价值）
        self.total_asset = self.amount + (self.stocks * price).sum()
        self.initial_total_asset = self.total_asset  # 记录初始总资产
        self.gamma_reward = 0.0  # 重置折扣奖励
        return self.get_state(price)  # 返回初始状态 , {}

    def step(self, actions):
        """
        执行一个动作（交易决策）

        参数:
        actions: 动作数组，表示每只股票的交易数量

        返回:
        state: 新状态
        reward: 奖励值
        done: 是否结束回合
        truncated: 是否被截断
        info: 额外信息
        """
        # 将动作缩放到实际交易数量
        actions = (actions * self.max_stock).astype(int)

        self.day += 1  # 进入下一个交易日
        price = self.price_ary[self.day]  # 获取当日股价
        self.stocks_cool_down += 1  # 所有股票冷却时间增加1

        # 检查当日是否处于高波动状态
        if self.turbulence_bool[self.day] == 0:
            # 正常交易状态
            min_action = int(self.max_stock * self.min_stock_rate)  # 计算最小交易量
            
            # 处理卖出操作（负动作值）
            for index in np.where(actions < -min_action)[0]:  # 找到需要卖出的股票索引
                if price[index] > 0:  # 确保股价有效
                    # 计算可卖出数量（持仓数量和动作数量的较小值）
                    sell_num_shares = min(self.stocks[index], -actions[index])
                    # 更新股票持仓
                    self.stocks[index] -= sell_num_shares
                    # 增加现金（考虑卖出成本）
                    self.amount += (
                        price[index] * sell_num_shares * (1 - self.sell_cost_pct)
                    )
                    # 重置该股票冷却时间
                    self.stocks_cool_down[index] = 0
                    
            # 处理买入操作（正动作值）
            for index in np.where(actions > min_action)[0]:  # 找到需要买入的股票索引
                if price[index] > 0:  # 确保股价有效
                    # 计算可买入数量（根据现金和动作数量确定）
                    buy_num_shares = min(self.amount // price[index], actions[index])
                    # 更新股票持仓
                    self.stocks[index] += buy_num_shares
                    # 减少现金（考虑买入成本）
                    self.amount -= (
                        price[index] * buy_num_shares * (1 + self.buy_cost_pct)
                    )
                    # 重置该股票冷却时间
                    self.stocks_cool_down[index] = 0

        else:  
            # 高波动状态：卖出所有股票
            self.amount += (self.stocks * price).sum() * (1 - self.sell_cost_pct)
            self.stocks[:] = 0  # 清空所有股票持仓
            self.stocks_cool_down[:] = 0  # 重置所有股票冷却时间

        # 获取新状态
        state = self.get_state(price)
        # 计算总资产
        total_asset = self.amount + (self.stocks * price).sum()
        # 计算奖励（资产变化量）
        reward = (total_asset - self.total_asset) * self.reward_scaling
        self.total_asset = total_asset  # 更新总资产

        # 计算折扣奖励
        self.gamma_reward = self.gamma_reward * self.gamma + reward
        # 检查是否达到最大步数
        done = self.day == self.max_step
        if done:
            # 回合结束时使用累积折扣奖励作为最终奖励
            reward = self.gamma_reward
            # 计算回合总回报率
            self.episode_return = total_asset / self.initial_total_asset

        return state, reward, done, dict()

    def get_state(self, price):
        """
        构造环境状态

        参数:
        price: 当前股价数组

        返回:
        state: 状态数组
        """
        # 对现金金额进行缩放
        amount = np.array(self.amount * (2**-12), dtype=np.float32)
        scale = np.array(2**-6, dtype=np.float32)  # 定义缩放因子
        # 拼接所有状态信息
        return np.hstack(
            (
                amount,  # 现金金额
                self.turbulence_ary[self.day],  # 当日波动率
                self.turbulence_bool[self.day],  # 当日波动率布尔值
                price * scale,  # 缩放后的股价
                self.stocks * scale,  # 缩放后的股票持仓
                self.stocks_cool_down,  # 股票冷却时间
                self.tech_ary[self.day],  # 当日技术指标
            )
        )  # state.astype(np.float32)

    @staticmethod
    def sigmoid_sign(ary, thresh):
        """
        应用带符号的sigmoid函数

        参数:
        ary: 输入数组
        thresh: 阈值

        返回:
        处理后的数组
        """
        def sigmoid(x):
            # 定义sigmoid函数，减去0.5使其输出以0为中心
            return 1 / (1 + np.exp(-x * np.e)) - 0.5

        # 应用sigmoid函数并乘以阈值
        return sigmoid(ary / thresh) * thresh